using System;
using System.Collections.Generic;

namespace LaplaceMeshDeformer
{
    /// <summary>
    /// 线性系统类
    /// 用于求解形如 Ax = b 的线性方程组
    /// </summary>
    public class LinearSystem
    {
        private SparseMatrix matrix; // 系统矩阵A
        private double[,] rightHandSide; // 右侧向量b
        private double[,] solution; // 解向量x
        private int size; // 系统大小

        /// <summary>
        /// 构造函数
        /// </summary>
        /// <param name="size">系统大小</param>
        public LinearSystem(int size)
        {
            this.size = size;
            this.matrix = new SparseMatrix(size, size);
            this.rightHandSide = new double[size, 3]; // 3D位置，包含x,y,z三个分量
            this.solution = new double[size, 3];
        }

        /// <summary>
        /// 设置矩阵元素
        /// </summary>
        /// <param name="row">行索引</param>
        /// <param name="col">列索引</param>
        /// <param name="value">元素值</param>
        public void SetMatrixElement(int row, int col, double value)
        {
            matrix.SetElement(row, col, value);
        }

        /// <summary>
        /// 获取矩阵元素
        /// </summary>
        /// <param name="row">行索引</param>
        /// <param name="col">列索引</param>
        /// <returns>元素值</returns>
        public double GetMatrixElement(int row, int col)
        {
            return matrix.GetElement(row, col);
        }

        /// <summary>
        /// 设置右侧向量元素
        /// </summary>
        /// <param name="row">行索引</param>
        /// <param name="component">分量索引 (0:x, 1:y, 2:z)</param>
        /// <param name="value">元素值</param>
        public void SetRightHandSide(int row, int component, double value)
        {
            if (component < 0 || component >= 3)
            {
                throw new ArgumentOutOfRangeException(nameof(component), "分量索引必须在0-2之间");
            }
            rightHandSide[row, component] = value;
        }

        /// <summary>
        /// 获取右侧向量元素
        /// </summary>
        /// <param name="row">行索引</param>
        /// <param name="component">分量索引 (0:x, 1:y, 2:z)</param>
        /// <returns>元素值</returns>
        public double GetRightHandSide(int row, int component)
        {
            if (component < 0 || component >= 3)
            {
                throw new ArgumentOutOfRangeException(nameof(component), "分量索引必须在0-2之间");
            }
            return rightHandSide[row, component];
        }

        /// <summary>
        /// 获取解向量元素
        /// </summary>
        /// <param name="row">行索引</param>
        /// <param name="component">分量索引 (0:x, 1:y, 2:z)</param>
        /// <returns>元素值</returns>
        public double GetSolution(int row, int component)
        {
            if (component < 0 || component >= 3)
            {
                throw new ArgumentOutOfRangeException(nameof(component), "分量索引必须在0-2之间");
            }
            return solution[row, component];
        }

        /// <summary>
        /// 求解线性系统 Ax = b
        /// 使用共轭梯度法求解
        /// </summary>
        /// <returns>是否成功求解</returns>
        public bool Solve()
        {
            // 为每个分量(x,y,z)求解一次线性系统
            for (int component = 0; component < 3; component++)
            {
                // 提取当前分量的右侧向量
                double[] b = new double[size];
                for (int i = 0; i < size; i++)
                {
                    b[i] = rightHandSide[i, component];
                }

                // 求解该分量的线性系统
                double[] x = new double[size];
                bool success = SolveConjugateGradient(b, x, Config.LinearSolver.MaxIterations, Config.LinearSolver.Tolerance);
                
                if (!success)
                {
                    return false;
                }

                // 保存解向量
                for (int i = 0; i < size; i++)
                {
                    solution[i, component] = x[i];
                }
            }

            return true;
        }

        /// <summary>
        /// 使用共轭梯度法求解线性系统 Ax = b
        /// </summary>
        /// <param name="b">右侧向量</param>
        /// <param name="x">解向量（输出）</param>
        /// <param name="maxIterations">最大迭代次数</param>
        /// <param name="tolerance">收敛容差</param>
        /// <returns>是否成功求解</returns>
        private bool SolveConjugateGradient(double[] b, double[] x, int maxIterations, double tolerance)
        {
            int n = b.Length;
            
            // 初始化解向量为0
            for (int i = 0; i < n; i++)
            {
                x[i] = 0.0;
            }

            // r = b - Ax
            double[] r = new double[n];
            double[] Ax = new double[n];
            matrix.MultiplyVector(x, Ax);
            
            for (int i = 0; i < n; i++)
            {
                r[i] = b[i] - Ax[i];
            }

            // p = r
            double[] p = new double[n];
            for (int i = 0; i < n; i++)
            {
                p[i] = r[i];
            }

            // r_dot_r = r·r
            double r_dot_r = 0.0;
            for (int i = 0; i < n; i++)
            {
                r_dot_r += r[i] * r[i];
            }

            double initial_r_dot_r = r_dot_r;

            // 迭代求解
            for (int iter = 0; iter < maxIterations; iter++)
            {
                if (Math.Sqrt(r_dot_r / initial_r_dot_r) < tolerance)
                {
                    return true; // 已收敛
                }

                // Ap = A·p
                double[] Ap = new double[n];
                matrix.MultiplyVector(p, Ap);

                // p_dot_Ap = p·Ap
                double p_dot_Ap = 0.0;
                for (int i = 0; i < n; i++)
                {
                    p_dot_Ap += p[i] * Ap[i];
                }

                // alpha = r·r / p·Ap
                double alpha = r_dot_r / p_dot_Ap;

                // x = x + alpha·p
                for (int i = 0; i < n; i++)
                {
                    x[i] += alpha * p[i];
                }

                // r = r - alpha·Ap
                for (int i = 0; i < n; i++)
                {
                    r[i] -= alpha * Ap[i];
                }

                // new_r_dot_r = r·r
                double new_r_dot_r = 0.0;
                for (int i = 0; i < n; i++)
                {
                    new_r_dot_r += r[i] * r[i];
                }

                // beta = new_r_dot_r / r_dot_r
                double beta = new_r_dot_r / r_dot_r;

                // p = r + beta·p
                for (int i = 0; i < n; i++)
                {
                    p[i] = r[i] + beta * p[i];
                }

                r_dot_r = new_r_dot_r;
            }

            // 达到最大迭代次数但未收敛
            return Math.Sqrt(r_dot_r / initial_r_dot_r) < tolerance * 10; // 放宽容差
        }

        /// <summary>
        /// 雅可比迭代法求解线性系统
        /// </summary>
        /// <param name="maxIterations">最大迭代次数</param>
        /// <param name="tolerance">收敛容差</param>
        /// <returns>求解是否成功收敛</returns>
        public bool SolveJacobi(int maxIterations = 1000, double tolerance = 1e-8)
        {
            int n = matrix.Rows;
            int rhsColumns = rightHandSide.GetLength(1);
            
            // 对每一个右侧向量求解
            for (int col = 0; col < rhsColumns; col++)
            {
                // 提取当前右侧向量
                double[] b = new double[n];
                for (int i = 0; i < n; i++)
                {
                    b[i] = rightHandSide[i, col];
                }
                
                // 初始化解向量
                double[] x = new double[n];
                double[] xNew = new double[n];
                
                // 迭代求解
                for (int iter = 0; iter < maxIterations; iter++)
                {
                    // 对每一行计算新的解
                    for (int i = 0; i < n; i++)
                    {
                        double sum = 0.0;
                        
                        // 计算非对角元素的贡献
                        for (int j = 0; j < n; j++)
                        {
                            if (i != j)
                            {
                                sum += matrix.GetElement(i, j) * x[j];
                            }
                        }
                        
                        // 对角元素
                        double aii = matrix.GetElement(i, i);
                        if (Math.Abs(aii) < 1e-10)
                        {
                            // 对角元素接近零，无法继续
                            return false;
                        }
                        
                        // 更新解
                        xNew[i] = (b[i] - sum) / aii;
                    }
                    
                    // 计算新旧解之间的差异
                    double maxDiff = 0.0;
                    for (int i = 0; i < n; i++)
                    {
                        maxDiff = Math.Max(maxDiff, Math.Abs(xNew[i] - x[i]));
                        x[i] = xNew[i];
                    }
                    
                    // 检查收敛条件
                    if (maxDiff < tolerance)
                    {
                        // 保存解
                        for (int i = 0; i < n; i++)
                        {
                            solution[i, col] = x[i];
                        }
                        break;
                    }
                    
                    // 如果达到最大迭代次数，仍保存当前解
                    if (iter == maxIterations - 1)
                    {
                        for (int i = 0; i < n; i++)
                        {
                            solution[i, col] = x[i];
                        }
                    }
                }
            }
            
            return true;
        }
        
        /// <summary>
        /// 高斯-赛德尔迭代法求解线性系统
        /// </summary>
        /// <param name="maxIterations">最大迭代次数</param>
        /// <param name="tolerance">收敛容差</param>
        /// <returns>求解是否成功收敛</returns>
        public bool SolveGaussSeidel(int maxIterations = 1000, double tolerance = 1e-8)
        {
            int n = matrix.Rows;
            int rhsColumns = rightHandSide.GetLength(1);
            
            // 对每一个右侧向量求解
            for (int col = 0; col < rhsColumns; col++)
            {
                // 提取当前右侧向量
                double[] b = new double[n];
                for (int i = 0; i < n; i++)
                {
                    b[i] = rightHandSide[i, col];
                }
                
                // 初始化解向量
                double[] x = new double[n];
                
                // 迭代求解
                for (int iter = 0; iter < maxIterations; iter++)
                {
                    // 记录最大变化量
                    double maxDiff = 0.0;
                    
                    // 对每一行计算新的解
                    for (int i = 0; i < n; i++)
                    {
                        double oldX = x[i];
                        double sum1 = 0.0; // 已更新变量的和
                        double sum2 = 0.0; // 未更新变量的和
                        
                        // 计算已更新变量(第0到i-1行)的贡献
                        for (int j = 0; j < i; j++)
                        {
                            sum1 += matrix.GetElement(i, j) * x[j];
                        }
                        
                        // 计算未更新变量(第i+1到n-1行)的贡献
                        for (int j = i + 1; j < n; j++)
                        {
                            sum2 += matrix.GetElement(i, j) * x[j];
                        }
                        
                        // 对角元素
                        double aii = matrix.GetElement(i, i);
                        if (Math.Abs(aii) < 1e-10)
                        {
                            // 对角元素接近零，无法继续
                            return false;
                        }
                        
                        // 更新解
                        x[i] = (b[i] - sum1 - sum2) / aii;
                        
                        // 更新最大变化量
                        maxDiff = Math.Max(maxDiff, Math.Abs(x[i] - oldX));
                    }
                    
                    // 检查收敛条件
                    if (maxDiff < tolerance)
                    {
                        // 保存解
                        for (int i = 0; i < n; i++)
                        {
                            solution[i, col] = x[i];
                        }
                        break;
                    }
                    
                    // 如果达到最大迭代次数，仍保存当前解
                    if (iter == maxIterations - 1)
                    {
                        for (int i = 0; i < n; i++)
                        {
                            solution[i, col] = x[i];
                        }
                    }
                }
            }
            
            return true;
        }
        
        /// <summary>
        /// 计算向量点积
        /// </summary>
        private static double Dot(double[] a, double[] b)
        {
            double result = 0.0;
            for (int i = 0; i < a.Length; i++)
            {
                result += a[i] * b[i];
            }
            return result;
        }
        
        /// <summary>
        /// 将系数矩阵和右侧向量转换为增广矩阵形式的字符串
        /// </summary>
        public override string ToString()
        {
            int n = matrix.Rows;
            int m = rightHandSide.GetLength(1);
            
            string result = $"线性系统 {n}x{n}, 右侧向量数: {m}\n";
            
            // 打印增广矩阵的前几行作为示例
            int displayRows = Math.Min(n, 5);
            for (int i = 0; i < displayRows; i++)
            {
                // 打印矩阵行
                for (int j = 0; j < Math.Min(n, 5); j++)
                {
                    result += $"{matrix.GetElement(i, j):F4} ";
                }
                
                if (n > 5)
                {
                    result += "... ";
                }
                
                // 打印右侧向量
                result += "| ";
                for (int j = 0; j < Math.Min(m, 3); j++)
                {
                    result += $"{rightHandSide[i, j]:F4} ";
                }
                
                if (m > 3)
                {
                    result += "...";
                }
                
                result += "\n";
            }
            
            if (n > 5)
            {
                result += "...\n";
            }
            
            return result;
        }
    }
}